from __future__ import with_statement
from khronos import *
from khronos.extras.components import *

class Worker(Atomic):
    state = "idle"
    def status(self):
        return self.state
        
    @chain
    def initializer(self):
        tools = self.get_root()["tools"]
        rng = self.get_rng()
        for day in ("Monday", "Tuesday", "Wednesday", "Thursday", "Friday"):
            with self.broadcasting(day):
                for shift, (start, end) in (("morning",   ( 8*HOUR, 12*HOUR)), 
                                            ("afternoon", (13*HOUR, 17*HOUR))):
                    # calculate the start and end dates for the next shift
                    now = self.get_timedelta()
                    shift_start = now.day() + start + rng.gauss(0, 5) * MINUTE
                    shift_end   = now.day() +   end + rng.gauss(0, 5) * MINUTE
                    if shift_start < now:
                        shift_start += DAY
                        shift_end   += DAY
                        
                    # wait until shift start
                    yield shift_start - now
                    # shift block
                    with self.broadcasting(shift):
                        yield self.work(tools, until=shift_end)
                # wait till midnight before broadcasting the next day
                yield DAY - self.get_timedelta().time()
        self.terminate()
        
    def work(self, tools, until):
        if isinstance(tools, Resource): return self.work_rsc(tools, until)
        elif isinstance(tools, Level):  return self.work_lvl(tools, until)
        elif isinstance(tools, Store):  return self.work_str(tools, until)
        raise TypeError("invalid tools type")
        
    @chain
    def work_rsc(self, tools, until):
        self.state = "acquiring tools"
        yield tools.acquire() % (until - self.get_timedelta())
        with self.broadcasting("job"):
            self.state = "working"
            yield Delay(self.get_rng().gauss(60, 15) * MINUTE) % (until - self.get_timedelta())
            del self.state
        tools.release()
        
    @chain
    def work_lvl(self, tools, until):
        self.state = "acquiring tools"
        yield tools.get(1) % (until - self.get_timedelta())
        with self.broadcasting("job"):
            self.state = "working"
            yield Delay(self.get_rng().gauss(60, 15) * MINUTE) % (until - self.get_timedelta())
        self.state = "releasing tools"
        yield tools.put(1)
        del self.state
        
    @chain
    def work_str(self, tools, until):
        self.state = "acquiring tools"
        t = (yield tools.get() % (until - self.get_timedelta())).result
        with self.broadcasting("job"):
            self.state = "working with %s" % (t,)
            yield Delay(self.get_rng().gauss(60, 15) * MINUTE) % (until - self.get_timedelta())
        self.state = "releasing %s" % (t,)
        yield tools.put(t)
        del self.state
        
class Workers(Composite):
    def initializer(self):
        for x in xrange(10):
            self.add(Worker(str(x)))
            
def model():
    root = Composite("root")
    root.add(Workers("workers"))
    #root.add(Resource(capacity=5, name="tools"))
    #root.add(Level(capacity=5, initial_level=5, name="tools"))
    root.add(Store(capacity=5, name="tools", 
                   initial_content=["hammer", "saw", "ladder", "paint", "crowbar"]))
    return root
    
def main():
    m = model()
    k = Khronos(m)
    k.single_run(DAY)
    
